#include "yolodetector.h"
#include "source/logger.h"
#include <QDebug>
#include <iostream>

YoloDetector::YoloDetector(QObject *parent) :
    QObject(parent), ctx(0), modelData(nullptr), modelLoaded(false)
{
    boxConfThreshold = 0.5;
    nmsThreshold = 0.45;
}

YoloDetector::~YoloDetector()
{
    if (modelLoaded) {
        rknn_destroy(ctx);
        if (modelData) {
            free(modelData);
        }
    }
}

void YoloDetector::dump_tensor_attr(const rknn_tensor_attr* attr) const
{
    std::string shape_str = attr->n_dims < 1 ? "" : std::to_string(attr->dims[0]);
    for (int i = 1; i < attr->n_dims; ++i) {
        shape_str += ", " + std::to_string(attr->dims[i]);
    }

    LogDebug << "  index=" << attr->index
             << ", name=" << attr->name
             << ", n_dims=" << attr->n_dims
             << ", dims=[" << shape_str.c_str() << "]"
             << ", n_elems=" << attr->n_elems
             << ", size=" << attr->size
             << ", w_stride=" << attr->w_stride
             << ", size_with_stride=" << attr->size_with_stride
             << ", fmt=" << get_format_string(attr->fmt)
             << ", type=" << get_type_string(attr->type)
             << ", qnt_type=" << get_qnt_type_string(attr->qnt_type)
             << ", zp=" << attr->zp
             << ", scale=" << attr->scale;
}

bool YoloDetector::loadModel(const QString &modelPath)
{
    int modelDataSize = 0;
    FILE *fp = fopen(modelPath.toStdString().c_str(), "rb");
    if (fp == nullptr) {
        LogWarning << "Open file" << modelPath << "----failed.";
        return false;
    }

    fseek(fp, 0, SEEK_END);
    modelDataSize = ftell(fp);
    fseek(fp, 0, SEEK_SET);

    modelData = (unsigned char*)malloc(modelDataSize);
    if (modelData == nullptr) {
        LogWarning << "Malloc model data failed.";
        fclose(fp);
        return false;
    }
    fread(modelData, 1, modelDataSize, fp);
    fclose(fp);
    LogDebug << "read loadModel Successfully modelDataSize:" << modelDataSize;

    int ret = rknn_init(&ctx, modelData, modelDataSize, 0, nullptr);
    if (ret < 0) {
        LogWarning << "rknn_init error :" << ret;
        free(modelData);
        return false;
    }
    rknn_sdk_version version;
    ret = rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &version, sizeof(rknn_sdk_version));
    if (ret < 0) {
        LogWarning << "rknn_init error ret: " << ret;
        return -1;
    }
    LogDebug << "sdk version: "<< version.api_version << " driver version: "<< version.drv_version;

    ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &ioNum, sizeof(ioNum));
    if (ret != 0) {
        LogWarning << "rknn_query RKNN_QUERY_IN_OUT_NUM error" << ret;
        return -1;
    }
    LogDebug << "model input num: "<< ioNum.n_input << " output num: "<< ioNum.n_output;


    rknn_tensor_attr inputAttrs[ioNum.n_input];
    memset(inputAttrs, 0, sizeof(inputAttrs));
    for (int i = 0; i < ioNum.n_input; i++) {
        inputAttrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(inputAttrs[i]), sizeof(rknn_tensor_attr));
        if (ret != 0) {
            LogWarning << "rknn_query RKNN_QUERY_INPUT_ATTR error" << ret;
            return false;
        }
        dump_tensor_attr(&(inputAttrs[i]));
    }

    rknn_tensor_attr outputAttrs[ioNum.n_output];
    memset(outputAttrs, 0, sizeof(outputAttrs));
    for (int i = 0; i < ioNum.n_output; i++) {
        outputAttrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(outputAttrs[i]), sizeof(rknn_tensor_attr));
        if (ret != 0) {
            LogWarning << "rknn_query RKNN_QUERY_OUTPUT_ATTR error" << ret;
            return false;
        }
        dump_tensor_attr(&(outputAttrs[i]));
    }

    rknn_tensor_attr output_attrs[ioNum.n_output];
    memset(output_attrs, 0, sizeof(output_attrs));
    for (int i = 0; i < ioNum.n_output; i++) {
        output_attrs[i].index = i;
        ret                   = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr));
        dump_tensor_attr(&(output_attrs[i]));
    }


    if (inputAttrs[0].fmt == RKNN_TENSOR_NCHW) {
        imgChannel = inputAttrs[0].dims[1];
        imgHeight = inputAttrs[0].dims[2];
        imgWidth = inputAttrs[0].dims[3];
    } else {
        imgHeight = inputAttrs[0].dims[1];
        imgWidth = inputAttrs[0].dims[2];
        imgChannel = inputAttrs[0].dims[3];
    }

    modelInputWidth = imgWidth;
    modelInputHeight = imgHeight;
    modelLoaded = true;
    LogDebug << "loadModel Successed modelInputWidth:"
             << modelInputWidth
             <<"  modelInputHeight"
             << modelInputHeight;
    return true;
}

bool YoloDetector::preprocess(const cv::Mat &image, cv::Mat &resizedImage)
{
    cv::cvtColor(image, resizedImage, cv::COLOR_BGR2RGB);
    cv::resize(resizedImage, resizedImage, cv::Size(modelInputWidth, modelInputHeight));
    return true;
}

std::vector<cv::Rect> YoloDetector::detect(const cv::Mat &image)
{
    std::vector<cv::Rect> detectedBoxes;
    if (!modelLoaded) {
        LogWarning << "Model is not loaded.";
        return detectedBoxes;
    }

    cv::Mat resizedImage;
    preprocess(image, resizedImage);

    rknn_input inputs[1];
    memset(inputs, 0, sizeof(inputs));
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8;
    inputs[0].size = modelInputWidth * modelInputHeight * imgChannel;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].buf = resizedImage.data;

    int ret = rknn_inputs_set(ctx, ioNum.n_input, inputs);
    if (ret < 0) {
        LogWarning << "rknn_inputs_set error" << ret;
        return detectedBoxes;
    }

    ret = rknn_run(ctx, nullptr);
    if (ret < 0) {
        LogWarning << "rknn_run error" << ret;
        return detectedBoxes;
    }

    rknn_output outputs[ioNum.n_output];
    memset(outputs, 0, sizeof(outputs));
    for (int i = 0; i < 3; i++) {
        outputs[i].want_float = 0;
    }

    ret = rknn_outputs_get(ctx, ioNum.n_output, outputs, nullptr);
    if (ret < 0) {
        LogWarning << "rknn_outputs_get error" << ret;
        return detectedBoxes;
    }



    LogDebug << "detectedBoxes:" << detectedBoxes ;
    // detect_result_group_t detect_result_group;
    // std::vector<float> out_scales = {1.0, 1.0, 1.0};
    // std::vector<int32_t> out_zps = {0, 0, 0};
    // post_process((int8_t*)outputs[0].buf, (int8_t*)outputs[1].buf, (int8_t*)outputs[2].buf, modelInputHeight, modelInputWidth,
    //              boxConfThreshold, nmsThreshold, 1.0, 1.0, out_zps, out_scales, &detect_result_group);

    // for (int i = 0; i < detect_result_group.count; i++) {
    //     detect_result_t* det_result = &(detect_result_group.results[i]);
    //     cv::Rect box;
    //     box.x = det_result->box.left;
    //     box.y = det_result->box.top;
    //     box.width = det_result->box.right - det_result->box.left;
    //     box.height = det_result->box.bottom - det_result->box.top;
    //     detectedBoxes.push_back(box);
    // }

    rknn_outputs_release(ctx, 3, outputs);


    return detectedBoxes;
}
void testYoloDetector()
{
    // 实例化YoloDetector
    LogDebug <<"go to testYoloDetector";
    YoloDetector detector;

    // 加载模型
    if (!detector.loadModel("master/test/yolov5s-640-640.rknn")) {
        LogError << "Failed to load YOLO model" ;
        return;
    }

    // 读取测试图像
    cv::Mat image = cv::imread("master/test/bus.jpg");
    if (image.empty()) {
        LogError << "Failed to load test image" ;
        return;
    }
    LogDebug << "img width: "<< image.cols << " Height: " << image.rows;
    // 进行检测
    std::vector<cv::Rect> boxes = detector.detect(image);

    // 输出检测结果
    for (const cv::Rect &box : boxes) {
        LogDebug << "Detected box: ["
                  << "x: " << box.x << ", "
                  << "y: " << box.y << ", "
                  << "width: " << box.width << ", "
                  << "height: " << box.height << "]";
        cv::rectangle(image, box, cv::Scalar(0, 255, 0), 2);
    }

    // 显示结果图像
    cv::imwrite("out.jpg", image);
}
